package com.pearl.generator.mybatisplus;

import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.config.*;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.config.builder.Entity;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;
import freemarker.template.Configuration;
import freemarker.template.Template;
import freemarker.template.TemplateException;
import lombok.NoArgsConstructor;

import java.io.*;
import java.util.*;

/**
 * 增强模板引擎
 *
 * @author TangDan
 * @version 1.0
 * @since 2022/6/29
 */
@NoArgsConstructor
public class EnhanceFreemarkerTemplateEngine extends FreemarkerTemplateEngine {

    private Configuration configuration;

    public EnhanceFreemarkerTemplateEngine init(ConfigBuilder configBuilder) {
        this.configuration = new Configuration(Configuration.DEFAULT_INCOMPATIBLE_IMPROVEMENTS);
        this.configuration.setDefaultEncoding(ConstVal.UTF8);
        this.configuration.setClassForTemplateLoading(FreemarkerTemplateEngine.class, "/");
        return this;
    }

    @Override
    public void writer(Map<String, Object> objectMap, String templatePath,File outputFile) throws Exception {
        Template template = this.configuration.getTemplate(templatePath);
        FileOutputStream fileOutputStream = new FileOutputStream(outputFile);

        try {
            template.process(objectMap, new OutputStreamWriter(fileOutputStream, ConstVal.UTF8));
        } catch (Throwable var9) {
            try {
                fileOutputStream.close();
            } catch (Throwable var8) {
                var9.addSuppressed(var8);
            }

            throw var9;
        }
        fileOutputStream.close();
    }

    /**
     * 自定义POJO模版支持
     *
     * @param customFile 自定义模板文件
     * @param tableInfo  表信息
     * @param objectMap  模板信息
     */
    @Override
    protected void outputCustomFile(Map<String, String> customFile, TableInfo tableInfo, Map<String, Object> objectMap) {
        // 1. 获取名称
        ConfigBuilder configBuilder = (ConfigBuilder) objectMap.get("config"); // 所有配置
        Entity entity = configBuilder.getStrategyConfig().entity();
        String entityName = entity.getNameConvert().entityNameConvert(tableInfo); // 原实体类名，经过配置过滤等逻辑
        // 2. 获取包路径
        String pojoPackageName = configBuilder.getPackageConfig().getPackageInfo("Other");
        // 3. 循环所有需要输出的文件，eg: Query.java -> /templates/generator/pojo.java.ftl
        objectMap.put("pojoPackageName", pojoPackageName);
        customFile.forEach((key, value) -> {
            // POJO
            String suffixName = key.replace(".java", "");// 类名后缀，eg:Query
            String pojoPath = joinPath(configBuilder.getGlobalConfig().getOutputDir(), pojoPackageName + "." + suffixName.toLowerCase());
            String fileName = String.format(pojoPath + File.separator + entityName + "%s", key); // 生成文件名
            String pojoName = String.format(entityName + "%s", suffixName); // 对象名称
            objectMap.put("pojoName", pojoName);
            objectMap.put("pojo" + suffixName + "Name", pojoName);
            objectMap.put("pojoPackage", pojoPackageName+"."+suffixName.toLowerCase());
            this.outputFile(new File(fileName), objectMap, value, this.getConfigBuilder().getInjectionConfig().isFileOverride()); // 输出
        });
        // mapStruct
        Boolean enableMapStruct = (Boolean) objectMap.get("enableMapStruct");
        if (enableMapStruct){
            String key="MapStruct.java";
            String value="/templates/generator/mapStruct.java.ftl";
            String suffixName = key.replace(".java", "");// 类名后缀，eg:Query
            String pojoPath = joinPath(configBuilder.getGlobalConfig().getOutputDir(), pojoPackageName + "." + suffixName.toLowerCase());
            String fileName = String.format(pojoPath + File.separator + entityName + "%s", "MapStructMapper.java"); // 生成文件名
            String pojoName = String.format(entityName + "%s", suffixName+"Mapper"); // 对象名称
            objectMap.put("mapStructPackage", pojoPackageName+"."+suffixName.toLowerCase());
            objectMap.put("pojoName", pojoName);
            objectMap.put("pojo" + suffixName + "Name", pojoName);
            this.outputFile(new File(fileName), objectMap, value, this.getConfigBuilder().getInjectionConfig().isFileOverride()); // 输出
        }
    }

    /**
     * 自定义POJO模版支持
     *
     * @param customFile 自定义模板文件
     * @param tableInfo  表信息
     * @param objectMap  模板信息
     */
    protected Map<String, String> previewCustomFile(Map<String, String> customFile, TableInfo tableInfo, Map<String, Object> objectMap) {
        // 1. 获取名称
        Map<String, String> templateMap = new HashMap<>();
        ConfigBuilder configBuilder = (ConfigBuilder) objectMap.get("config"); // 所有配置
        Entity entity = configBuilder.getStrategyConfig().entity();
        String entityName = entity.getNameConvert().entityNameConvert(tableInfo); // 原实体类名，经过配置过滤等逻辑
        String pojoPackageName = configBuilder.getPackageConfig().getPackageInfo("Other");
        objectMap.put("pojoPackageName", pojoPackageName);
        // 2. 循环所有需要输出的文件，eg: Query.java -> /templates/generator/pojo.java.ftl
        customFile.forEach((key, value) -> {
            String suffixName = key.replace(".java", "");// VO
            String pojoName = String.format(entityName + "%s", suffixName); // 对象名称
            objectMap.put("pojoName", pojoName);
            objectMap.put("pojo" + suffixName + "Name", pojoName);// 不同类型不同变量，一个表多个一样模板时，会导致冲突
            objectMap.put("pojoPackage", pojoPackageName+"."+suffixName.toLowerCase());
            Map<String, String> pojoTemplateMap = this.previewFile(objectMap, value);
            String templateValue = pojoTemplateMap.get("/templates/generator/pojo.java.ftl");// POJO 都是使用的同一个模板
            String pojoKey = value.replace("pojo", suffixName); // pojo 替换为VO、DTO，前端分类预览
            Map<String, String> previewPojoTemplateMap = new HashMap<>();
            previewPojoTemplateMap.put(pojoKey, templateValue);
            templateMap.putAll(previewPojoTemplateMap);
        });
        // mapStruct
        Boolean enableMapStruct = (Boolean) objectMap.get("enableMapStruct");
        if (enableMapStruct){
            String key="MapStruct.java";
            String value="/templates/generator/mapStruct.java.ftl";
            String suffixName = key.replace(".java", "");// 类名后缀，eg:Query
            String pojoPath = joinPath(configBuilder.getGlobalConfig().getOutputDir(), pojoPackageName + "." + suffixName.toLowerCase());
            String fileName = String.format(pojoPath + File.separator + entityName + "%s", "MapStructMapper.java"); // 生成文件名
            String pojoName = String.format(entityName + "%s", suffixName+"Mapper"); // 对象名称
            objectMap.put("mapStructPackage", pojoPackageName+"."+suffixName.toLowerCase());
            objectMap.put("pojoName", pojoName);
            objectMap.put("pojo" + suffixName + "Name", pojoName);
            Map<String, String> pojoTemplateMap = this.previewFile(objectMap, value);
            String templateValue = pojoTemplateMap.get("/templates/generator/mapStruct.java.ftl");// POJO 都是使用的同一个模板
            String pojoKey = value.replace("pojo", suffixName); // pojo 替换为VO、DTO，前端分类预览
            Map<String, String> previewPojoTemplateMap = new HashMap<>();
            previewPojoTemplateMap.put(pojoKey, templateValue);
            templateMap.putAll(previewPojoTemplateMap);
        }
        return templateMap;
    }

    /**
     * 拼接输出文件路径
     *
     * @param parentDir
     * @param packageName
     * @return
     */
    private String joinPath(String parentDir, String packageName) {
        if (StringUtils.isBlank(parentDir)) {
            parentDir = System.getProperty("java.io.tmpdir");
        }
        if (!StringUtils.endsWith(parentDir, File.separator)) {
            parentDir = parentDir + File.separator;
        }
        packageName = packageName.replaceAll("\\.", "\\" + File.separator);
        return parentDir + packageName;
    }

    /**
     * 自定义注入模板变量
     *
     * @param config    配置
     * @param tableInfo 表信息
     * @return 模板变量
     */
    @Override
    public Map<String, Object> getObjectMap(ConfigBuilder config, TableInfo tableInfo) {
        StrategyConfig strategyConfig = config.getStrategyConfig();
        Map<String, Object> controllerData = strategyConfig.controller().renderData(tableInfo);
        Map<String, Object> objectMap = new HashMap<>(controllerData);
        Map<String, Object> mapperData = strategyConfig.mapper().renderData(tableInfo);
        objectMap.putAll(mapperData);
        Map<String, Object> serviceData = strategyConfig.service().renderData(tableInfo);
        objectMap.putAll(serviceData);
        Map<String, Object> entityData = strategyConfig.entity().renderData(tableInfo);
        objectMap.putAll(entityData);
        objectMap.put("config", config);
        objectMap.put("package", config.getPackageConfig().getPackageInfo());
        GlobalConfig globalConfig = config.getGlobalConfig();
        objectMap.put("author", globalConfig.getAuthor());
        objectMap.put("kotlin", globalConfig.isKotlin());
        objectMap.put("swagger", globalConfig.isSwagger());
        // 设置全局Lombok 模式
        objectMap.put("lombok", (Boolean) Objects.requireNonNull(config.getInjectionConfig()).getCustomMap().get("lombok"));
        objectMap.put("date", globalConfig.getCommentDate());
        String schemaName = "";
        if (strategyConfig.isEnableSchema()) {
            schemaName = config.getDataSourceConfig().getSchemaName();
            if (StringUtils.isNotBlank(schemaName)) {
                schemaName = schemaName + ".";
                tableInfo.setConvert(true);
            }
        }
        objectMap.put("schemaName", schemaName);
        objectMap.put("table", tableInfo);
        objectMap.put("entity", tableInfo.getEntityName());
        return objectMap;
    }

    /**
     * 生成预览文件
     *
     * @param objectMap    变量
     * @param templatePath 模板路径
     * @return 生成的预览
     */
    protected Map<String, String> previewFile(Map<String, Object> objectMap, String templatePath) {
        try {
            // 获取模板
            // "/templates/generator/BO.java.ftl
            Template template = this.configuration.getTemplate(templatePath);
            StringWriter stringWriter = new StringWriter();
            BufferedWriter writer = new BufferedWriter(stringWriter);
            // 处理
            template.process(objectMap, writer);
            StringReader reader = new StringReader(stringWriter.toString());
            writer.flush();
            writer.close();
            BufferedReader r = new BufferedReader(reader);
            StringBuilder sb = new StringBuilder();
            String line;
            while ((line = r.readLine()) != null) {
                sb.append(line);
                sb.append("\r\n");
            }
            Map<String, String> mapper = new HashMap<>();
            mapper.put(templatePath, sb.toString());
            return mapper;
        } catch (IOException | TemplateException e) {
            throw new RuntimeException(e);
        }
    }




    /**
     * 获取预览模板
     *
     * @return 模板
     */
    public Map<String, String> preview() {
        try {
            Map<String, String> allTemplateMap = new HashMap<String, String>();
            ConfigBuilder config = this.getConfigBuilder();
            List<TableInfo> tableInfoList = config.getTableInfoList();
            for (TableInfo tableInfo : tableInfoList) {
                Map<String, Object> objectMap = this.getObjectMap(config, tableInfo);
                // POJO
                Optional.ofNullable(config.getInjectionConfig()).ifPresent((t) -> {
                    t.beforeOutputFile(tableInfo, objectMap);
                    allTemplateMap.putAll(this.previewCustomFile(t.getCustomFile(), tableInfo, objectMap));
                });
                // Mapper
                allTemplateMap.putAll(this.previewMapper(tableInfo, objectMap));
                // Entity
                allTemplateMap.putAll(this.previewEntity(tableInfo, objectMap));
                // Service
                allTemplateMap.putAll(this.previewService(tableInfo, objectMap));
                // Controller
                allTemplateMap.putAll(this.previewController(tableInfo, objectMap));
            }
            return allTemplateMap;
        } catch (Exception var3) {
            throw new RuntimeException("无法创建文件，请检查配置信息！", var3);
        }
    }

    /**
     * 获取Service层预览模板
     *
     * @param tableInfo 表信息
     * @param objectMap 变量
     * @return 模板
     */
    protected Map<String, String> previewService(TableInfo tableInfo, Map<String, Object> objectMap) {
        Map<String, String> templateMap = new HashMap<>();
        String servicePath = this.getPathInfo(OutputFile.service);
        if (StringUtils.isNotBlank(tableInfo.getServiceName()) && StringUtils.isNotBlank(servicePath)) {
            this.getTemplateFilePath(TemplateConfig::getService).ifPresent((service) -> {
                Map<String, String> serviceMap = previewFile(objectMap, service);
                templateMap.putAll(serviceMap);
            });
        }
        String serviceImplPath = this.getPathInfo(OutputFile.serviceImpl);
        if (StringUtils.isNotBlank(tableInfo.getServiceImplName()) && StringUtils.isNotBlank(serviceImplPath)) {
            this.getTemplateFilePath(TemplateConfig::getServiceImpl).ifPresent((serviceImpl) -> {
                Map<String, String> serviceImplMap = previewFile(objectMap, serviceImpl);
                templateMap.putAll(serviceImplMap);
            });
        }
        return templateMap;
    }

    /**
     * 获取Controller层预览模板
     *
     * @param tableInfo 表信息
     * @param objectMap 变量
     * @return 模板
     */
    protected Map<String, String> previewController(TableInfo tableInfo, Map<String, Object> objectMap) {
        String controllerPath = this.getPathInfo(OutputFile.controller);
        Map<String, String> templateMap = new HashMap<>();
        if (StringUtils.isNotBlank(tableInfo.getControllerName()) && StringUtils.isNotBlank(controllerPath)) {
            this.getTemplateFilePath(TemplateConfig::getController).ifPresent((controller) -> {
                Map<String, String> controllerMap = this.previewFile(objectMap, controller);
                templateMap.putAll(controllerMap);
            });
        }
        return templateMap;
    }

    /**
     * 获取Entity层预览模板
     *
     * @param tableInfo 表信息
     * @param objectMap 变量
     * @return 模板
     */
    protected Map<String, String> previewEntity(TableInfo tableInfo, Map<String, Object> objectMap) {
        String entityName = tableInfo.getEntityName();
        String entityPath = this.getPathInfo(OutputFile.entity);
        Map<String, String> templateMap = new HashMap<>();
        if (StringUtils.isNotBlank(entityName) && StringUtils.isNotBlank(entityPath)) {
            Optional<String> templateFilePath = this.getTemplateFilePath((template) -> template.getEntity(this.getConfigBuilder().getGlobalConfig().isKotlin()));
            templateFilePath.ifPresent((entity) -> {
                Map<String, String> entityMap = this.previewFile(objectMap, entity);
                templateMap.putAll(entityMap);
            });
        }
        return templateMap;
    }

    /**
     * 获取Mapper层预览模板
     *
     * @param tableInfo 表信息
     * @param objectMap 变量
     * @return 模板
     */
    protected Map<String, String> previewMapper(TableInfo tableInfo, Map<String, Object> objectMap) {
        Map<String, String> templateMap = new HashMap<>();
        String xmlPath = this.getPathInfo(OutputFile.xml);
        if (StringUtils.isNotBlank(tableInfo.getXmlName()) && StringUtils.isNotBlank(xmlPath)) {
            this.getTemplateFilePath(TemplateConfig::getXml).ifPresent((xml) -> {
                Map<String, String> xmlMap = previewFile(objectMap, xml);
                templateMap.putAll(xmlMap);
            });
        }
        String mapperPath = this.getPathInfo(OutputFile.mapper);
        if (StringUtils.isNotBlank(tableInfo.getMapperName()) && StringUtils.isNotBlank(mapperPath)) {
            this.getTemplateFilePath(TemplateConfig::getMapper).ifPresent((mapper) -> {
                Map<String, String> mapperMap = previewFile(objectMap, mapper);
                templateMap.putAll(mapperMap);
            });
        }
        return templateMap;
    }

}
